package io.github.javpower.knowledgerag.core;

import java.util.ArrayList;
import java.util.List;

public class HybridBlockSplitter {

    public List<String> split(String document, int blockSize) {
        List<String> blocks = new ArrayList<>();
        List<String> tokens = Tokenizer.tokenize(document);

        StringBuilder currentBlock = new StringBuilder();
        int currentBlockSize = 0;

        for (String token : tokens) {
            if (currentBlockSize + token.length() > blockSize) {
                if (currentBlock.length() > 0) {
                    blocks.add(currentBlock.toString().trim());
                    currentBlock = new StringBuilder();
                    currentBlockSize = 0;
                }
                // 如果 token 本身超过块大小，进行递归分块
                if (token.length() > blockSize) {
                    List<String> subBlocks = splitRecursive(token, blockSize);
                    blocks.addAll(subBlocks);
                } else {
                    currentBlock.append(token).append(" ");
                    currentBlockSize += token.length();
                }
            } else {
                currentBlock.append(token).append(" ");
                currentBlockSize += token.length();
            }
        }

        if (currentBlock.length() > 0) {
            blocks.add(currentBlock.toString().trim());
        }

        return blocks;
    }

    private List<String> splitRecursive(String text, int blockSize) {
        List<String> blocks = new ArrayList<>();
        int start = 0;
        while (start < text.length()) {
            int end = Math.min(start + blockSize, text.length());
            int splitPoint = text.lastIndexOf(' ', end);
            if (splitPoint == -1 || splitPoint == start) {
                splitPoint = end;
            }
            String block = text.substring(start, splitPoint).trim();
            blocks.add(block);
            start = splitPoint;
        }
        return blocks;
    }
}